# -*- coding: utf-8 -*-

import numpy as np
import cv2 as cv
import pandas as pd

import argparse
import glob
import sys
import logging

import core

parser = argparse.ArgumentParser(description='This sample demonstrates Lucas-Kanade Optical Flow calculation.')
parser.add_argument('video', type=str, help='path to video file')
args = None
if len(sys.argv) > 1:
    args = parser.parse_args()
else:
    args = parser.parse_args(glob.glob("../data/*.avi")[0:1])

# params for ShiTomasi corner detection
feature_params = dict(maxCorners=1000,
                      qualityLevel=0.3,
                      minDistance=3,
                      blockSize=6,
                      # useHarrisDetector=True,
                      # k=0.005
                      )

# Parameters for lucas kanade optical flow
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT,
                           10, 0.03)
                 )

# 视频间隔的帧数，最小为0
# 比如设定为4，将处理第5、10、15、20...帧，即间隔4帧
FRAMES_INTERVAL = 4
# 视频播放一帧的时间, 无限长为"NAN"
ONE_FRAME_PLAY_TIME = 30

# Create some random colors
color = np.random.randint(0, 255, (1000, 3))

# 帧数记忆, 帧遍历时累加和赋值, 不要再做额外的修改
frame_count = 0
loop_count = 0
# 保存到文件的光流数据
df = None

cap = cv.VideoCapture(args.video)
# Take first frame and find corners in it
ret, old_frame = cap.read()
# old_frame = cv.transpose(old_frame)
old_gray = cv.cvtColor(old_frame, cv.COLOR_BGR2GRAY)
p0 = cv.goodFeaturesToTrack(old_gray, mask=None, **feature_params)
logging.info("p0.type: " + str(type(p0)))
logging.info("p0.dim: " + str(p0.shape))

# Create a mask image for drawing purposes
mask = np.zeros_like(old_frame)

while True:

    ret, frame = cap.read()

    frame_count += 1
    loop_count += 1
    if frame_count <= FRAMES_INTERVAL:
        continue
    else:
        frame_count = 0

    # frame = cv.transpose(frame)
    if frame is None:
        break
    frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)

    # calculate optical flow
    p1, st, err = cv.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)

    # Select good points
    good_new = p1[st == 1]
    good_old = p0[st == 1]
    if df is None:
        df = pd.DataFrame(np.hstack((good_old, good_new)),
                          columns=pd.MultiIndex.from_product(
                              [[1], ["old", "new"], ['x', 'y']]
                          ))

    temp_df = pd.DataFrame(
        np.hstack((good_old, good_new)),
        columns=pd.MultiIndex.from_product(
            [[loop_count], ["old", "new"], ['x', 'y']]
        ))
    df = pd.concat([df, temp_df], axis=1)
    # draw the tracks
    for i, (new, old) in enumerate(zip(good_new, good_old)):
        a, b = new.ravel()
        c, d = old.ravel()
        mask = cv.line(mask, (a, b), (c, d), color[i].tolist(), 2)
        frame = cv.circle(frame, (a, b), 5, color[i].tolist(), -1)
    img = cv.add(frame, mask)

    cv.imshow('frame', img)
    k = None
    if ONE_FRAME_PLAY_TIME == "NAN":
        k = cv.waitKey() & 0xff
    else:
        k = cv.waitKey(ONE_FRAME_PLAY_TIME) & 0xff
    if k == 27:
        break

    # Now update the previous frame and previous points
    old_gray = frame_gray.copy()
    # p0 = good_new.reshape(-1, 1, 2)
    p0 = cv.goodFeaturesToTrack(old_gray, mask=None, **feature_params)


frame1 = old_frame
prvs = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
hsv = np.zeros_like(frame1)
hsv[..., 1] = 255
cv.destroyAllWindows()
df.to_excel("./optical_flow.xlsx")